import numpy as np
from scipy import rand
from scipy.optimize import leastsq
import pylab

class HypObj:
     """Object for best fit to hyperbolic function"""
     def __init__(self,x,y,guess=None,plot=None):
         self.x = x
         self.y = y
         self.plot=plot
         if guess==None:
             self.guess = [y.max(),1.0,0.5]
         else:
             self.guess = guess
         self.parameters, self.success = leastsq(self._errf,
                                                 self.guess[:],
                                                 args = (self.x, self.y))
         self.r2 = self._r2(self.x,self.y,self.parameters)
         if self.plot<>None:
             pylab.plot(self.x, self.y, 'ro',
                        self.x, self._f(self.parameters, self.x), 'b--')
             pylab.show()
             pylab.clf()
     def _f(self,p, x):
         """Evaluate function with parameters, p, and array of x values"""
         return p[0]*(1.0+p[1]*p[2]*x)**(-1.0/p[1])
     def _errf(self,p, x, y):
         """return the difference between calculated and input y values"""
         return self._f(p, x) - y
     def _r2(self,x,y,abc):
         """calculate the correlation coefficient and rmsd"""
         y_calc = self._f(abc,self.x)
         y_avg = y.mean()
         rmsd = ((y-self._f(abc,x))**2).sum()/((y_calc-y_avg)**2).sum()
         return (1.0-rmsd, rmsd)

def list_toarray(data):
     """Convert input list of data pairs to two arrays"""
     x = []
     y = []
     for i in data:
         x.append(float(i[0]))
         y.append(float(i[1]))
     x = np.array(x)
     y = np.array(y)
     return x,y

def test():

     def f(p, x): # for testing, make the objective function available
         return p[0]*(1.0+p[1]*p[2]*x)**(-1.0/p[1])

     # fake data for testing, should get perfect correlation
     print "\nTest 1: should find an exact solution = [1.25,0.75,0.25]"
     print "(No plot)"
     qi, b, di = [1.25,0.75,0.25]
     sample = []
     for i in range(1,61):
         sample.append([float(i),qi*(1.0+b*di*float(i))**(-1.0/b)])
     x, y = list_toarray(sample)
     ex1 = HypObj(x,y)
     if ex1.success==1:
         print ex1.parameters,ex1.r2
     else:
         print "leastsq() error: ",ex1.success

     # generate data with random deviations
     print "\nTest 2: should find a solution close to [15.0, 2.5, 0.3]"
     print "(No plot)"
     n = 151
     x = np.mgrid[1:10:n*1j]
     abc = [15.0, 2.5, 0.3]
     y = f(abc, x) + rand(n) - 0.5
     ex2 = HypObj(x,y,guess=[30.0, 4.0, 1.0])
     if ex2.success==1:
         print ex2.parameters,ex2.r2
     else:
         print "leastsq() error: ",ex2.success

     # real data
     print "\nTest 3: real world data"
     rn115604=[[1, 3233],[2, 3530],[3, 3152],[4, 2088],[6, 3038],
               [7, 2108],[8, 2132],[9, 1654],[10, 1762],[11, 1967],
               [12, 1760],[13, 1649],[14, 1633],[15, 1680],[16, 1398],
               [17, 1622],[18, 1393],[19, 1436],[20, 1352],[21, 1402],
               [22, 1459],[23, 1373],[24, 1262],[25, 1346],[26, 1325],
               [27, 1319],[28, 1309],[29, 1206],[30, 1249],[31, 1446],
               [32, 1255],[33, 1227],[34, 1268],[35, 1233],[36, 1175],
               [37, 1129],[38, 1242],[39, 1247],[40, 1198],[41, 1058],
               [42, 1172],[43, 1242],[44, 1214],[45, 1148],[46, 1689],
               [47, 971],[48, 1084],[49, 1028],[50, 1164],[51, 1297],
               [52, 1040],[53, 1045],[54, 1196],[55, 991],[56, 1065],
               [57, 898],[58, 1020],[59, 966],[60, 1162],[61, 1069],
               [62, 1055],[63, 1035],[64, 1045],[65, 1076],[66, 1108],
               [67, 918],[68, 1051],[69, 1049],[70, 1039],[71, 1133],
               [72, 887],[73, 924],[74, 983],[75, 1077],[76, 1092],
               [77, 973],[78, 920],[79, 1040]]
     x, y = list_toarray(rn115604)
     ex3 = HypObj(x,y,plot=1)
     if ex3.success==1:
         print ex3.parameters,ex3.r2
     else:
         print "leastsq() error: ",ex3.success

if __name__=="__main__":
     test()
